from ultralytics import SETTINGS
from ultralytics.models import YOLO
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

if __name__ == '__main__':
	model = YOLO(model='ultralytics/cfg/models/11/yolo11-use.yaml')
	# model.load('yolov8n.pt')
	SETTINGS["tensorboard"] = True
	model.train(data='./data.yaml', epochs=150, batch=4, device='0', imgsz=640, workers=0, cache=False,
				amp=True, mosaic=False, project='runs/train', name='exp', exist_ok=False,
				# 启动tensorboard
				plots=True,  # 启用训练曲线图
				save_period=5,  # 每5个epoch保存一次模型
				)	# exist_ok: 迭代实验
	SETTINGS["tensorboard"] = False